package state

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"

	"github.com/louloulin/dataflare/pkg/workflow/engine"
)

// StateStore represents a state store
type StateStore interface {
	// Get gets a value
	Get(ctx context.Context, key string) ([]byte, error)
	// Set sets a value
	Set(ctx context.Context, key string, value []byte) error
	// Delete deletes a value
	Delete(ctx context.Context, key string) error
	// List lists keys with a prefix
	List(ctx context.Context, prefix string) ([]string, error)
}

// StateManager represents a state manager
type StateManager struct {
	// store is the state store
	store StateStore
}

// NewStateManager creates a new state manager
func NewStateManager(store StateStore) *StateManager {
	return &StateManager{
		store: store,
	}
}

// SaveWorkflowState saves workflow state
func (m *StateManager) SaveWorkflowState(workflowID string, runID string, state *engine.WorkflowStatus) error {
	// Create context
	ctx := context.Background()

	// Marshal state
	data, err := json.Marshal(state)
	if err != nil {
		return fmt.Errorf("failed to marshal workflow state: %w", err)
	}

	// Save state
	key := fmt.Sprintf("workflow:%s:run:%s", workflowID, runID)
	err = m.store.Set(ctx, key, data)
	if err != nil {
		return fmt.Errorf("failed to save workflow state: %w", err)
	}

	return nil
}

// GetWorkflowState gets workflow state
func (m *StateManager) GetWorkflowState(workflowID string, runID string) (*engine.WorkflowStatus, error) {
	// Create context
	ctx := context.Background()

	// Get state
	key := fmt.Sprintf("workflow:%s:run:%s", workflowID, runID)
	data, err := m.store.Get(ctx, key)
	if err != nil {
		return nil, fmt.Errorf("failed to get workflow state: %w", err)
	}

	// Unmarshal state
	var state engine.WorkflowStatus
	err = json.Unmarshal(data, &state)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal workflow state: %w", err)
	}

	return &state, nil
}

// SaveNodeState saves node state
func (m *StateManager) SaveNodeState(workflowID string, runID string, nodeID string, state *engine.NodeStatus) error {
	// Create context
	ctx := context.Background()

	// Marshal state
	data, err := json.Marshal(state)
	if err != nil {
		return fmt.Errorf("failed to marshal node state: %w", err)
	}

	// Save state
	key := fmt.Sprintf("workflow:%s:run:%s:node:%s", workflowID, runID, nodeID)
	err = m.store.Set(ctx, key, data)
	if err != nil {
		return fmt.Errorf("failed to save node state: %w", err)
	}

	return nil
}

// GetNodeState gets node state
func (m *StateManager) GetNodeState(workflowID string, runID string, nodeID string) (*engine.NodeStatus, error) {
	// Create context
	ctx := context.Background()

	// Get state
	key := fmt.Sprintf("workflow:%s:run:%s:node:%s", workflowID, runID, nodeID)
	data, err := m.store.Get(ctx, key)
	if err != nil {
		return nil, fmt.Errorf("failed to get node state: %w", err)
	}

	// Unmarshal state
	var state engine.NodeStatus
	err = json.Unmarshal(data, &state)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal node state: %w", err)
	}

	return &state, nil
}

// ListRunsByWorkflow lists all workflow runs
func (m *StateManager) ListRunsByWorkflow(workflowID string) ([]string, error) {
	// Create context
	ctx := context.Background()

	// List runs
	prefix := fmt.Sprintf("workflow:%s:run:", workflowID)
	keys, err := m.store.List(ctx, prefix)
	if err != nil {
		return nil, fmt.Errorf("failed to list workflow runs: %w", err)
	}

	// Extract run IDs
	runIDs := make([]string, 0, len(keys))
	for _, key := range keys {
		// Extract run ID from key
		runID := key[len(prefix):]
		runIDs = append(runIDs, runID)
	}

	return runIDs, nil
}

// InMemoryStateStore represents an in-memory state store
type InMemoryStateStore struct {
	// data is the in-memory data
	data map[string][]byte
	// mu is the mutex
	mu sync.RWMutex
}

// NewInMemoryStateStore creates a new in-memory state store
func NewInMemoryStateStore() *InMemoryStateStore {
	return &InMemoryStateStore{
		data: make(map[string][]byte),
	}
}

// Get gets a value
func (s *InMemoryStateStore) Get(ctx context.Context, key string) ([]byte, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	// Check if the key exists
	value, ok := s.data[key]
	if !ok {
		return nil, fmt.Errorf("key %s not found", key)
	}

	return value, nil
}

// Set sets a value
func (s *InMemoryStateStore) Set(ctx context.Context, key string, value []byte) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Set the value
	s.data[key] = value

	return nil
}

// Delete deletes a value
func (s *InMemoryStateStore) Delete(ctx context.Context, key string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Delete the key
	delete(s.data, key)

	return nil
}

// List lists keys with a prefix
func (s *InMemoryStateStore) List(ctx context.Context, prefix string) ([]string, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	// Find keys with the prefix
	keys := make([]string, 0)
	for key := range s.data {
		if len(key) >= len(prefix) && key[:len(prefix)] == prefix {
			keys = append(keys, key)
		}
	}

	return keys, nil
}

// FileStateStore represents a file-based state store
type FileStateStore struct {
	// TODO: Implement file-based state store
}

// RedisStateStore represents a Redis-based state store
type RedisStateStore struct {
	// TODO: Implement Redis-based state store
}

// PostgresStateStore represents a PostgreSQL-based state store
type PostgresStateStore struct {
	// TODO: Implement PostgreSQL-based state store
}

// NewStateStore creates a new state store based on the type
func NewStateStore(storeType string, config map[string]interface{}) (StateStore, error) {
	switch storeType {
	case "memory":
		return NewInMemoryStateStore(), nil
	case "file":
		return nil, fmt.Errorf("file state store not implemented")
	case "redis":
		return nil, fmt.Errorf("redis state store not implemented")
	case "postgres":
		return nil, fmt.Errorf("postgres state store not implemented")
	default:
		return nil, fmt.Errorf("unsupported state store type: %s", storeType)
	}
}
